package com.lfy.generator.utils;


import com.baomidou.mybatisplus.core.toolkit.StringPool;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.baomidou.mybatisplus.generator.config.ConstVal;
import com.baomidou.mybatisplus.generator.config.builder.ConfigBuilder;
import com.baomidou.mybatisplus.generator.engine.AbstractTemplateEngine;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.io.IOUtils;
import org.apache.velocity.Template;
import org.apache.velocity.VelocityContext;
import org.apache.velocity.app.Velocity;
import org.apache.velocity.app.VelocityEngine;

import java.io.ByteArrayOutputStream;
import java.io.StringWriter;
import java.util.Map;
import java.util.Properties;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

/**
 * @Author: 李丰翼
 * @DateTime: 2020/11/4 0004 16:23
 * @Description: 自定义输出模版配置
 */
@Slf4j
public class MyVelocityTemplateEngine extends AbstractTemplateEngine {
    private static final String DOT_VM = ".vm";
    private VelocityEngine velocityEngine;
    private ZipOutputStream zip;
    private ByteArrayOutputStream outputStream;

    public ZipOutputStream getZip() {
        return zip;
    }

    public ByteArrayOutputStream getOutputStream() {
        return outputStream;
    }

    @SneakyThrows
    @Override
    public MyVelocityTemplateEngine init(ConfigBuilder configBuilder) {
        super.init(configBuilder);
        if (null == velocityEngine) {
            Properties p = new Properties();
            p.setProperty(ConstVal.VM_LOAD_PATH_KEY, ConstVal.VM_LOAD_PATH_VALUE);
            p.setProperty(Velocity.FILE_RESOURCE_LOADER_PATH, StringPool.EMPTY);
            p.setProperty(Velocity.ENCODING_DEFAULT, ConstVal.UTF8);
            p.setProperty(Velocity.INPUT_ENCODING, ConstVal.UTF8);
            p.setProperty("file.resource.loader.unicode", StringPool.TRUE);
            outputStream = new ByteArrayOutputStream();
            zip = new ZipOutputStream(outputStream);
//            zip = new ZipOutputStream(new FileOutputStream("D:\\2.zip"));
            velocityEngine = new VelocityEngine(p);
        }
        return this;
    }


    @Override
    public void writer(Map<String, Object> objectMap, String templatePath, String outputFile) throws Exception {
        if (StringUtils.isBlank(templatePath)) {
            return;
        }
        Template template = velocityEngine.getTemplate(templatePath, ConstVal.UTF8);
        try (
                StringWriter sw = new StringWriter()
        ) {
            template.merge(new VelocityContext(objectMap), sw);


            //添加到zip
            zip.putNextEntry(new ZipEntry(outputFile));
            IOUtils.write(sw.toString(), zip, "UTF-8");
            IOUtils.closeQuietly(sw);
            zip.closeEntry();
            log.info("模板:" + templatePath + ";  文件:" + outputFile);
        }

    }


    @Override
    public String templateFilePath(String filePath) {
        if (null == filePath || filePath.contains(DOT_VM)) {
            return filePath;
        }
        return filePath + DOT_VM;
    }
}
